import os


def read_video_names(datalist_path):
    video_names = []
    with open(datalist_path, "r") as f:
        for line in f:
            content = line.strip()
            if len(content) == 0:
                continue
            video_names.append(content)
    return video_names


def prepare_HOI4D_Sim_data(BundleTrack_data_category_dir, video_names, output_dir):
    for video_name in video_names:
        print(video_name)
        source_dir = os.path.join(BundleTrack_data_category_dir, video_name.replace("/", "_"))
        target_dir = os.path.join(output_dir, video_name)
        os.makedirs(os.path.join(target_dir, "rgb"), exist_ok=True)
        os.makedirs(os.path.join(target_dir, "depth"), exist_ok=True)

        for i in range(300):
            os.system("cp {} {}".format(os.path.join(source_dir, "rgb", str(i).zfill(4) + ".png"), os.path.join(target_dir, "rgb", str(i).zfill(5) + ".png")))
            os.system("cp {} {}".format(os.path.join(source_dir, "depth", str(i).zfill(4) + ".png"), os.path.join(target_dir, "depth", str(i).zfill(5) + ".png")))


if __name__ == "__main__":
    categories = [
        "C2",
        "C5",
        "C7",
        "C12",
    ]
    datalist_dir = "/home/liuyun/HOI4D-Sim-Vision-Pipeline/datalists"
    dataset_dir = "/localdata_sdd/HOI4D_Sim"
    BundleTrack_data_root_dir = "/localdata_sdd/HOI4D_BundleTrack"

    for category in categories:
        video_names = read_video_names(os.path.join(datalist_dir, category + ".txt"))
        BundleTrack_data_category_dir = os.path.join(BundleTrack_data_root_dir, category)
        prepare_HOI4D_Sim_data(BundleTrack_data_category_dir, video_names, dataset_dir)
